//
//  APDUCommands.swift
//  NFCPassportReader
//
//  Created by Manwel Bugeja Personal on 16/04/2025.
//
//

import Foundation
import CoreNFC
import CryptoKit
import Security
import OSLog

private let logger = Logger(subsystem: Bundle.main.bundleIdentifier ?? "com.example.nfcreader", category: "APDUCommands")

/**
 * Defines the type of cryptographic key to be used for a signature operation.
 *
 * Use this enum to specify whether an operation should use the card's
 * authentication key or its non-repudiable signing key.
 */
public enum SignatureKeyType: CustomStringConvertible, CaseIterable, Identifiable {
    /// The key pair used for authentication purposes.
    case authentication
    /// The key pair used for non-repudiable digital signatures.
    case signing
    
    public var description: String {
        switch self {
        case .authentication: return "Authentication"
        case .signing: return "Signing"
        }
    }
    
    // Conformance for use in SwiftUI Pickers
    public static var allCases: [SignatureKeyType] {
        return [.authentication, .signing]
    }
    public var id: Self { self }
}

/**
 * Specifies the signature padding scheme to be requested from the eID card.
 *
 * This enum defines the supported signature formats and their corresponding
 * algorithm reference bytes for the MSE:SET APDU command.
 */
public enum SignatureFormat: CustomStringConvertible, CaseIterable, Identifiable {
    /// Represents the RSASSA-PKCS1-v1_5 signature scheme.
    case pkcs1v15
    /// Represents the RSASSA-PSS signature scheme.
    case rsaPSS
    
    /// The algorithm reference byte used in the MSE:SET APDU command.
    var algorithmReference: UInt8 {
        switch self {
        case .pkcs1v15: return 0xE8 // Algorithm reference for PKCS#1 v1.5
        case .rsaPSS:   return 0xEC // Algorithm reference for RSA-PSS
        }
    }
    
    public var description: String {
        switch self {
        case .pkcs1v15: return "PKCS#1 v1.5"
        case .rsaPSS: return "RSA-PSS"
        }
    }
    
    // Conformance for use in SwiftUI Pickers
    public static var allCases: [SignatureFormat] {
        return [.pkcs1v15, .rsaPSS]
    }
    public var id: Self { self }
}

/// A utility enum for sending APDU commands to a Maltese eID card.
///
/// This enum acts as a namespace for orchestrating low-level APDU commands required
/// for cryptographic operations such as reading digital certificates and generating
/// digital signatures. It provides a high-level API to simplify interaction with the card.
///
/// The public methods assume that a `TagReader` with an active, secure session
/// (e.g., after PACE authentication) is provided.
public enum APDUCommands {
    
    // MARK: - Private Constants
    
    /// The Application Identifier for the Master File (MF) of the eID card.
    private static let MF_AID: [UInt8] = [0xD2, 0x76, 0x00, 0x00, 0x98, 0x4D, 0x4C, 0x41, 0x56, 0x31]
    /// The Dedicated File (DF) Identifier for authentication-related objects.
    private static let DF_AUTH_ID: [UInt8] = [0x50, 0x15]
    /// The Dedicated File (DF) Identifier for signing-related objects.
    private static let DF_SIG_ID: [UInt8] = [0x1F, 0xFF]
    /// The File Identifier (FID) for the authentication certificate.
    private static let AUTH_CERT_FID: [UInt8] = [0x54, 0x02]
    /// The File Identifier (FID) for the signing certificate.
    private static let SIGNATURE_CERT_FID: [UInt8] = [0x54, 0x01]

    // MARK: - Public API
    
    /**
     Performs a complete digital signature process using either the authentication or signing key.
     
     This function handles the entire APDU command sequence required to generate a signature:
     1. Selects the Master File.
     2. Selects the appropriate Directory File (DF) for the given key type.
     3. Verifies the user PIN for the given key type.
     4. Sets the security environment (MSE:SET) for the digital signature operation.
     5. Sends the challenge to the card to compute the digital signature.
     
     - Note: The `challenge` data is sent directly to the card. If the selected signature algorithm
       (e.g., one specified by `customAlgorithmReference: 0xEC`) requires a pre-hashed input, the caller
       is responsible for performing that hash *before* calling this function.
     
     - Parameters:
        - keyType: The key to use for signing, either `.authentication` or `.signing`.
        - tagReader: An authenticated `TagReader` instance with an active secure session.
        - pin: The PIN corresponding to the selected `keyType`.
        - challenge: The data to be signed (which may be a pre-computed hash).
        - format: The desired signature format, which determines the algorithm reference if `customAlgorithmReference` is `nil`. Defaults to `.rsaPSS`.
        - customAlgorithmReference: An optional byte to override the algorithm reference from `format`.
                                   For the Maltese eID card, `0xEC` is used to sign a pre-computed hash.
     - Returns: A byte array (`[UInt8]`) containing the computed digital signature from the card.
     - Throws: `NFCPassportReaderError` if any step in the communication fails, such as an incorrect PIN, a communication error, or an unexpected response from the card.
    */
    public static func performSignature(
        using keyType: SignatureKeyType,
        on tagReader: TagReader,
        pin: String,
        challenge: Data,
        format: SignatureFormat = .rsaPSS,
        customAlgorithmReference: UInt8? = nil
    ) async throws -> [UInt8] {
        
        try await selectMasterFile(tagReader: tagReader)
        try await selectDF(for: keyType, tagReader: tagReader)
        try await verifyPIN(for: keyType, tagReader: tagReader, pin: pin)
        
        try await setupSecurityEnvironment(
            for: keyType,
            tagReader: tagReader,
            signatureFormat: format,
            customAlgorithmReference: customAlgorithmReference
        )
        
        let signatureResponse = try await computeDigitalSignature(tagReader: tagReader, challenge: challenge)

        guard signatureResponse.sw1 == 0x90, signatureResponse.sw2 == 0x00 else {
            let algoRefStr = customAlgorithmReference.map { String(format: "0x%02X", $0) } ?? "Default (\(format))"
            throw NFCPassportReaderError.ResponseError(
                "Failed to compute digital signature for \(keyType) key (AlgoRef: \(algoRefStr))",
                signatureResponse.sw1,
                signatureResponse.sw2
            )
        }
        
        return signatureResponse.data
    }
    
    /**
     Reads the raw X.509 certificate data from the eID card for the specified key type.
     
     This function selects the appropriate files on the card and reads the binary data
     of the certificate.
     
     - Parameters:
        - tagReader: An authenticated `TagReader` instance with an active secure session.
        - keyType: The key type (`.authentication` or `.signing`) whose certificate should be read.
     - Returns: A byte array (`[UInt8]`) containing the raw certificate data.
     - Throws: `NFCPassportReaderError` if the certificate cannot be selected or read.
     */
    public static func readCertificate(tagReader: TagReader, keyType: SignatureKeyType) async throws -> [UInt8] {
        try await selectMasterFile(tagReader: tagReader)
        try await selectDF(for: keyType, tagReader: tagReader)
        
        let fileId = (keyType == .authentication) ? AUTH_CERT_FID : SIGNATURE_CERT_FID
        try await selectCertificateFile(tagReader: tagReader, fileId: fileId)
        
        return try await readBinaryFile(tagReader: tagReader)
    }
    
    /**
     Parses raw byte data into a `SecCertificate` object.
     
     This is a utility function to convert the data read from the card into a usable
     certificate object from the `Security` framework.
     
     - Parameter certData: The raw byte array of the X.509 certificate.
     - Returns: A `SecCertificate` object.
     - Throws: `NFCPassportReaderError.InvalidDataPassed` if the data cannot be parsed.
     */
    public static func parseCertificate(certData: [UInt8]) throws -> SecCertificate {
        guard let certificate = SecCertificateCreateWithData(nil, Data(certData) as CFData) else {
            throw NFCPassportReaderError.InvalidDataPassed("Failed to parse certificate data")
        }
        return certificate
    }
    
    /**
     Extracts the public key (`SecKey`) from a `SecCertificate` object.
     
     - Parameter certificate: The certificate from which to extract the public key.
     - Returns: The extracted `SecKey` public key.
     - Throws: `NFCPassportReaderError.InvalidDataPassed` if the public key cannot be extracted.
     */
    public static func extractPublicKey(certificate: SecCertificate) throws -> SecKey {
        guard let publicKey = SecCertificateCopyKey(certificate) else {
            throw NFCPassportReaderError.InvalidDataPassed("Failed to extract public key from certificate")
        }
        return publicKey
    }

    // MARK: - Deprecated Functions (for backward compatibility)

    @available(*, deprecated, renamed: "performSignature(using:on:pin:challenge:format:customAlgorithmReference:)", message: "Use the new unified `performSignature` method instead.")
    public static func performCompleteAuthenticationKeySignature(
        tagReader: TagReader,
        pin: String,
        challenge: Data,
        signatureFormat: SignatureFormat = .rsaPSS,
        customAlgorithmReference: UInt8? = nil
    ) async throws -> [UInt8] {
        try await performSignature(
            using: .authentication,
            on: tagReader,
            pin: pin,
            challenge: challenge,
            format: signatureFormat,
            customAlgorithmReference: customAlgorithmReference
        )
    }
    
    @available(*, deprecated, renamed: "performSignature(using:on:pin:challenge:format:customAlgorithmReference:)", message: "Use the new unified `performSignature` method instead.")
    public static func performCompleteSigningKeySignature(
        tagReader: TagReader,
        pin: String,
        challenge: Data,
        signatureFormat: SignatureFormat = .rsaPSS,
        customAlgorithmReference: UInt8? = nil
    ) async throws -> [UInt8] {
        try await performSignature(
            using: .signing,
            on: tagReader,
            pin: pin,
            challenge: challenge,
            format: signatureFormat,
            customAlgorithmReference: customAlgorithmReference
        )
    }

    // MARK: - Private APDU Command Helpers
    
    /// Sends a command to select the Master File (root) of the card's file system.
    private static func selectMasterFile(tagReader: TagReader) async throws {
        let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0xA4, p1Parameter: 0x04, p2Parameter: 0x0C, data: Data(MF_AID), expectedResponseLength: -1)
        let response = try await tagReader.send(cmd: cmd)
        guard response.sw1 == 0x90, response.sw2 == 0x00 else {
            throw NFCPassportReaderError.ResponseError("Failed to select Master File (MF)", response.sw1, response.sw2)
        }
    }
    
    /// Sends a command to select a Dedicated File (DF) for either Authentication or Signing.
    private static func selectDF(for keyType: SignatureKeyType, tagReader: TagReader) async throws {
        let dfId = (keyType == .authentication) ? DF_AUTH_ID : DF_SIG_ID
        let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0xA4, p1Parameter: 0x00, p2Parameter: 0x0C, data: Data(dfId), expectedResponseLength: -1)
        let response = try await tagReader.send(cmd: cmd)
        guard response.sw1 == 0x90, response.sw2 == 0x00 else {
            throw NFCPassportReaderError.ResponseError("Failed to select \(keyType) DF", response.sw1, response.sw2)
        }
    }

    /// Sends a command to verify the PIN for a specific key type (Authentication or Signing).
    private static func verifyPIN(for keyType: SignatureKeyType, tagReader: TagReader, pin: String) async throws {
        guard let pinData = pin.data(using: .isoLatin1) else {
            throw NFCPassportReaderError.InvalidDataPassed("PIN contains non ISO-8859-1 characters.")
        }
        
        let p2: UInt8 = (keyType == .authentication) ? 0x83 : 0x92
        let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0x20, p1Parameter: 0x00, p2Parameter: p2, data: pinData, expectedResponseLength: -1)
        
        let response = try await tagReader.send(cmd: cmd)
        // Note: Specific error codes like 63CX (retries left) or 6983 (blocked) could be handled here for more detailed errors.
        guard response.sw1 == 0x90, response.sw2 == 0x00 else {
            throw NFCPassportReaderError.ResponseError("\(keyType) PIN verification failed", response.sw1, response.sw2)
        }
    }

    /// Sends a command to set the security environment (MSE:SET) for a cryptographic operation.
    /// This tells the card which key and algorithm to use for the subsequent command.
    private static func setupSecurityEnvironment(
        for keyType: SignatureKeyType,
        tagReader: TagReader,
        signatureFormat: SignatureFormat,
        customAlgorithmReference: UInt8?
    ) async throws {
        let algoRef = customAlgorithmReference ?? signatureFormat.algorithmReference
        let keyRef: UInt8 = (keyType == .authentication) ? 0x42 : 0x41
        
        // MSE:SET data: Tag '80' (Algo) + Len 1 + [algoRef] | Tag '84' (Key) + Len 1 + [keyRef]
        let mseData: [UInt8] = [0x80, 0x01, algoRef, 0x84, 0x01, keyRef]
        let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0x22, p1Parameter: 0x41, p2Parameter: 0xB6, data: Data(mseData), expectedResponseLength: -1)
        
        logger.debug("Sending MSE:SET (\(keyType)) with AlgoRef: \(String(format: "0x%02X", algoRef)), KeyRef: \(String(format: "0x%02X", keyRef))")
        let response = try await tagReader.send(cmd: cmd)
        
        guard response.sw1 == 0x90, response.sw2 == 0x00 else {
            throw NFCPassportReaderError.ResponseError(
                "Failed to set security environment for \(keyType) (AlgoRef: \(String(format: "0x%02X", algoRef)))",
                response.sw1,
                response.sw2
            )
        }
    }
    
    /// Sends the command to compute the digital signature (PSO:CDS) over the provided challenge.
    private static func computeDigitalSignature(tagReader: TagReader, challenge: Data) async throws -> ResponseAPDU {
        let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0x2A, p1Parameter: 0x9E, p2Parameter: 0x9A, data: challenge, expectedResponseLength: 256)
        return try await tagReader.send(cmd: cmd, useExtendedMode: true)
    }
    
    /// Sends a command to select an Elementary File (EF), specifically a certificate file.
    private static func selectCertificateFile(tagReader: TagReader, fileId: [UInt8]) async throws {
        let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0xA4, p1Parameter: 0x02, p2Parameter: 0x0C, data: Data(fileId), expectedResponseLength: -1)
        let response = try await tagReader.send(cmd: cmd)
        guard response.sw1 == 0x90, response.sw2 == 0x00 else {
            throw NFCPassportReaderError.ResponseError("Failed to select certificate file", response.sw1, response.sw2)
        }
    }
    
    /// Reads the full contents of the currently selected binary file (EF).
    /// This function handles reading in chunks until the end of the file is reached.
    private static func readBinaryFile(tagReader: TagReader) async throws -> [UInt8] {
        var allData = [UInt8]()
        var offset = 0
        let chunkSize = 200 // A reasonable chunk size to avoid overwhelming the card buffer.
        
        while true {
            let offsetHigh = UInt8((offset >> 8) & 0xFF)
            let offsetLow = UInt8(offset & 0xFF)
            
            let cmd = NFCISO7816APDU(instructionClass: 0x0C, instructionCode: 0xB0, p1Parameter: offsetHigh, p2Parameter: offsetLow, data: Data(), expectedResponseLength: chunkSize)
            let response = try await tagReader.send(cmd: cmd)
            
            // Success with data is the primary condition.
            if response.sw1 == 0x90, response.sw2 == 0x00 {
                if response.data.isEmpty { break } // No more data to read.
                
                allData.append(contentsOf: response.data)
                offset += response.data.count
                
                // If the returned data is less than what we asked for, it's the last chunk.
                if response.data.count < chunkSize { break }
            } else {
                // If we've already read some data, a different status might indicate a clean end-of-file.
                // If we haven't read anything yet, it's a genuine error.
                if allData.isEmpty {
                    throw NFCPassportReaderError.ResponseError("Failed to read binary data", response.sw1, response.sw2)
                } else {
                    // Assume end of file on any non-9000 status after successful reads have already occurred.
                    break
                }
            }
        }
        return allData
    }
}
